/*
 * Copyright (C) 2020 original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
 * in compliance with the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License
 * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 * or implied. See the License for the specific language governing permissions and limitations under
 * the License.
 */

package vip.justlive.easyproxy.server.repository;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.TypeReference;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Collection;
import java.util.Map;
import java.util.OptionalInt;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import vip.justlive.easyproxy.core.context.Config;
import vip.justlive.easyproxy.core.context.Metrics;
import vip.justlive.easyproxy.core.context.RemoteContext;
import vip.justlive.easyproxy.core.context.ServerProxyContext;
import vip.justlive.easyproxy.core.protocol.ProxyData;
import vip.justlive.easyproxy.server.entity.Client;
import vip.justlive.easyproxy.server.entity.Proxy;
import vip.justlive.easyproxy.server.handler.ConnectionUpdater;
import vip.justlive.easyproxy.server.handler.ProxyHandler;
import vip.justlive.easyproxy.server.handler.ServerHandler;
import vip.justlive.oxygen.core.config.ConfigFactory;
import vip.justlive.oxygen.core.config.Ignore;
import vip.justlive.oxygen.core.exception.Exceptions;
import vip.justlive.oxygen.core.net.aio.core.GroupContext;
import vip.justlive.oxygen.core.net.aio.core.Server;
import vip.justlive.oxygen.core.util.FileUtils;
import vip.justlive.oxygen.core.util.ThreadUtils;
import vip.justlive.oxygen.ioc.annotation.Bean;
import vip.justlive.oxygen.ioc.annotation.PostConstruct;
import vip.justlive.oxygen.ioc.annotation.PreDestroy;

/**
 * repository
 *
 * @author wubo
 */
@Bean
@Slf4j
@RequiredArgsConstructor
public class AdminRepository {

  private static final String FILE_NAME = "db.data";
  private static final AtomicInteger IDS = new AtomicInteger(1);

  private final Config config;
  private final ServerHandler handler;

  @Ignore
  private Server server;
  private Map<String, Client> clients = new ConcurrentHashMap<>();

  public void saveClient(String key, String name) {
    Client client = clients.get(key);
    if (client == null) {
      clients.put(key, new Client().setKey(key).setName(name));
    } else {
      client.setName(name);
    }
    store();
  }

  public Client getClient(String key) {
    return clients.get(key);
  }

  public Collection<Client> clients() {
    return clients.values();
  }

  public void deleteClient(String key) {
    Client client = clients.remove(key);
    if (client == null) {
      return;
    }
    store();
    for (Proxy proxy : client.getProxies()) {
      closeProxy(key, proxy.getId());
    }
    client.getProxies().clear();
    RemoteContext remote = RemoteContext.lookup(key);
    if (remote != null) {
      remote.getChannel().write(ProxyData.close());
    }
  }

  public void saveProxy(String key, Proxy proxy) {
    Client client = getClient(key);
    if (client == null) {
      throw Exceptions.fail("client not found");
    }
    Proxy local = client.find(proxy.getId());
    if (local == null) {
      local = client.find(proxy.getRemote(), proxy.getPort());
    }
    if (local != null) {
      local.setName(proxy.getName()).setRemote(proxy.getRemote()).setPort(proxy.getPort());
    } else {
      client.getProxies().add(proxy.setId(IDS.getAndIncrement()));
    }
    store();
    startProxy(key, proxy);
  }

  public void deleteProxy(String key, int id) {
    Client client = getClient(key);
    if (client == null) {
      return;
    }
    client.getProxies().removeIf(proxy -> proxy.getId() == id);
    store();
    closeProxy(key, id);
  }

  public void startProxy(String key, Proxy proxy) {
    RemoteContext context = RemoteContext.lookup(key);
    if (context == null) {
      return;
    }
    closeProxy(key, proxy.getId());
    ServerProxyContext spc = new ServerProxyContext().setKey(key).setName(proxy.getName())
        .setRemote(proxy.getRemote()).setPort(proxy.getPort()).setId(proxy.getId());
    context.getProxies().add(spc);
    createProxy(spc, context);
  }

  public void createProxy(ServerProxyContext proxy, RemoteContext remote) {
    ProxyHandler proxyHandler = new ProxyHandler(proxy, remote.getChannel());
    GroupContext ctx = new GroupContext(proxyHandler).setAioListener(proxyHandler);
    ConfigFactory.load(ctx, Config.AIO_PROXY_CONFIG);
    Server server = new Server(ctx);
    try {
      server.start(new InetSocketAddress(config.getServer().getHost(), proxy.getPort()));
    } catch (IOException e) {
      throw Exceptions.wrap(e);
    }
    proxy.setServer(server);
    Metrics.metrics(proxy.getPort());
  }

  @PostConstruct
  public void init() throws IOException {
    GroupContext context = new GroupContext(handler);
    ConfigFactory.load(context, Config.AIO_CONFIG);
    context.setAioListener(handler);
    server = new Server(context);
    server.start(config.getServer().getHost(), config.getServer().getPort());
    log.info("proxy server started {}", server.getServerAddress());

    Metrics.metrics(config.getServer().getPort());
    ThreadUtils.globalTimer()
        .scheduleWithFixedDelay(new ConnectionUpdater(context, this), 1, 1, TimeUnit.SECONDS);

    FileUtils.mkdirs(config.getAdmin().getStore());
    Path path = Paths.get(config.getAdmin().getStore(), FILE_NAME);
    if (!path.toFile().exists()) {
      return;
    }
    byte[] data = Files.readAllBytes(path);
    if (data.length > 0) {
      clients = JSON.parseObject(data, new TypeReference<TreeMap<String, Client>>() {
      }.getType());
      OptionalInt max = clients.values().stream()
          .flatMapToInt(client -> client.getProxies().stream().mapToInt(Proxy::getId)).max();
      if (max.isPresent()) {
        IDS.set(max.getAsInt() + 1);
      }
    }
  }

  @PreDestroy
  public void destroy() {
    if (server != null) {
      server.stop();
      log.info("close proxy server {}", server.getServerAddress());
    }
  }

  private void closeProxy(String key, int id) {
    RemoteContext context = RemoteContext.lookup(key);
    if (context == null) {
      return;
    }
    ServerProxyContext proxy = context.remove(id);
    if (proxy == null || proxy.getServer() == null) {
      return;
    }
    proxy.getServer().stop();
    context.getChannel().write(ProxyData.closeProxy(id));
  }


  private void store() {
    try {
      Files.write(Paths.get(config.getAdmin().getStore(), FILE_NAME), JSON.toJSONBytes(clients));
    } catch (IOException e) {
      log.error("store clients error", e);
    }
  }
}
